import os
import glob
import torch
import datetime
import gradio as gr
from diffusers import DiffusionPipeline
from diffusers import StableDiffusionPipeline
from diffusers import DDIMScheduler, DPMSolverMultistepScheduler, EulerAncestralDiscreteScheduler

device = "cuda" if torch.cuda.is_available() else "cpu"


# ui界面
def ui():
    from theme.summer import Summer
    with gr.Blocks(title="Text To Image Of Stable Diffusion", theme=Summer()) as interface:
        gr.Markdown('''# <span style='color:brown'>AI BOY VIP</span> ''')
        with gr.Row():
            with gr.Column():
                models = get_sd_model('./models', 'safetensors')
                lora_models = get_sd_model('./models/lora', 'safetensors')
                sd_model = gr.Dropdown(label="模型", choices=models, value=models[0], interactive=True)
                lora = gr.Dropdown(label="LORA", choices=lora_models, interactive=True)
                with gr.Row():
                    prompt = gr.Textbox(label="正向提示词", lines=10)
                    negative_prompt = gr.Textbox(label="反向提示词", lines=10)
                with gr.Row():
                    width = gr.Number(label="宽", value=512, step=2)
                    height = gr.Number(label="高", value=512, step=2)
                with gr.Row():
                    steps = gr.Number(label="采样步数", value=20, step=1)
                    guidance = gr.Number(label="提示词相关性", value=7.5, step=0.5)
                    num_images = gr.Number(label="每批数量", value=1, step=1)
                with gr.Row():
                    save_img = gr.Checkbox(label="保存图片", value=True)
                model_button = gr.Button(value="GO GO GO", variant='primary')
            with gr.Column():
                gallery_img = gr.Gallery(label="AIBoyVIP", visible=True, preview=True)
        model_button.click(txt_to_img,
                           inputs=[sd_model, lora, prompt, negative_prompt, width, height, steps, guidance, num_images,
                                   save_img], outputs=gallery_img)

        interface.queue(max_size=8).launch()


# 文生图功能
def txt_to_img(sd_model: str, lora: str, prompt: str, negative_prompt: str, width: str, height: str, steps: str,
               guidance: str,
               num_images: str, save_img: str):
    width = int(width)
    height = int(height)
    steps = int(steps)
    num_images = int(num_images)
    save_img = bool(save_img)
    guidance = float(guidance)

    pipeline = StableDiffusionPipeline.from_single_file(
        sd_model,
        torch_dtype=torch.float16,
        local_files_only=True,
        force_download=False,
        resume_download=False,
        load_safety_checker=False
    ).to(torch_device=device)

    # 切换为DPM采样器（这里应该可以通过ui改成自定义采样器）
    pipeline.scheduler = DDIMScheduler.from_config(pipeline.scheduler.config)
    pipeline.safety_checker = None
    pipeline.requires_safety_checker = False

    # 加载lora
    if lora:
        pipeline.load_lora_weights(lora, local_files_only=True, force_download=False)

    # 通过pipeline进行文生图
    images = pipeline(prompt=prompt, negative_prompt=negative_prompt, width=width, height=height,
                      num_inference_steps=steps, guidance_scale=guidance, num_images_per_prompt=num_images).images

    # 保存图片
    if save_img:
        for img in images:
            now = datetime.datetime.now()
            formatted_date = now.strftime('%Y%m%d%H%M%S%f')
            img.save('img/' + formatted_date + '.png')
    return images


# 获取指定路径下指定格式的所有文件
def get_sd_model(folder_path: str, file_extension: str):
    def get_files(folder_path, file_extension):
        # 构建文件搜索路径
        search_path = os.path.join(folder_path, f"*.{file_extension}")
        # 使用glob模块搜索匹配的文件路径
        matched_files = glob.glob(search_path)
        return matched_files

    # 获取匹配的文件列表
    files = get_files(folder_path, file_extension)

    models = []
    for file_path in files:
        models.append(file_path)
        print(f'加载模型 {file_path}')
    return models


if __name__ == '__main__':
    ui()
